"""Module for functions updating the primitive variables."""

import numpy

class IEPrimitiveVariable(object):
    def __init__(self, grid):
        self.grid = grid

    def update_primitive_variables(self):
        grid = self.grid
                
        w = grid.w; q = grid.q

        w[0][:] = q[0]
        w[1][:] = q[1]/q[0]

class EEPrimitiveVariable(object):
    def __init__(self, gamma, grid):
        self.grid = grid
        self.gamma = gamma

    def update_primitive_variables(self):

        grid = self.grid
        gamma = self.gamma

        # get the vector of conserved variables
        q = grid.q

        # get the vector of primitive variables
        w = grid.w

        rho = q[0]

        u = q[1]/rho
        e = q[2]/rho - 0.5 * u * u
        p = (gamma - 1.0) * rho * e

        w[0][:] = rho[:]
        w[1][:] = u[:]
        w[2][:] = p[:]

        # set the sound speed
        grid.a = numpy.sqrt( gamma * p/rho )

class EEPrimitiveVariableLagrangian(object):
    def __init__(self, gamma, grid):
        self.grid = grid
        self.gamma = gamma

    def update_primitive_variables(self):

        grid = self.grid
        gamma = self.gamma

        # get the vector of conserved variables
        q = grid.q

        # get the vector of primitive variables
        w = grid.w

        rho = q[0]

        u = q[1]/rho
        e = q[2]/rho - 0.5 * u * u
        p = (gamma - 1.0) * rho * e

        w[0][:] = rho[:]
        w[1][:] = u[:]
        w[2][:] = p[:]

        # set the sound speed
        grid.a = numpy.sqrt( gamma * p*rho )

